require 'utils'
require 'cutils'
require 'dict'

function create_data() 
	-- load dic 
	print('load dic')
	local f = torch.DiskFile(dic_emb_path, 'r')
	local dic = f:readObject()
	setmetatable(dic, Dict_mt)
	f:close()

	-- load wfreq --
	print('load word count')
	prob_to_select = {}
	w_prob = torch.zeros(dic:size())
	total_cases = 0
	ntgwords = 0
	nmajorwords = 0

	for line in io.lines(wfreq_path) do
		local comp = split_string(line)
		if #comp ~= 2 then error('invalid word freq file') end

		local wid = tonumber(comp[1]+1)
		local freq = tonumber(comp[2])

		prob_to_select[wid] = math.min(1, num_cases / freq)
		w_prob[{wid}] = freq

		total_cases = total_cases + freq
		ntgwords = ntgwords + 1
		if prob_to_select[wid] < 1 then 
			nmajorwords = nmajorwords + 1
		end
	end

	print('total cases: ' .. total_cases)
	print('total target words: ' .. ntgwords)
	print('total major words: ' .. nmajorwords)
	print('expected cases: ' .. num_cases * ntgwords)

	-- load data --
	print('load data')
	local nctxwords = 0
	for line in io.lines(ctx_data_path) do
		if line:len() > 0 then
			local comp = split_string(line)
			nctxwords = #comp - 1
			break
		end
	end

	local f = torch.DiskFile(ctx_data_path, 'r')
	local total_data = 	torch.IntTensor(f:readInt((nctxwords+1)*total_cases))
						:resize(total_cases, nctxwords+1):add(1)
	f:close()

	-- extract data, just take maximum num_cases per words
	print('draw data')
	local data = torch.IntTensor(1.5*num_cases*ntgwords, nctxwords + 1)
	local next_id = 1
	for i = 1,total_cases do
		local datum = total_data[{i,{}}]
		local wid = datum[{1}]
		local p = prob_to_select[wid]
		if p > math.random() then 
			data[{next_id,{}}]:copy(datum)
			next_id = next_id + 1
		end
	end
	data = data[{{1,next_id-1},{}}]
	data = data:t()
	data = shuffle(data)
	print(data:size())

	-- save to file
	print('save to file')
	f = torch.DiskFile(output_path, 'w')
	f:binary()
	f:writeObject(data)
	f:close()
end

if #arg ~= 5 then 
	print(
		"<ctx_data_path> " ..
		"<wfreq_path> " ..
		"<dic_emb_path> " ..
		"<number_cases_per_word> " ..
		"<output_path>")
else
	ctx_data_path	= arg[1]
	wfreq_path		= arg[2]
	dic_emb_path	= arg[3]
	num_cases		= tonumber(arg[4])
	output_path		= arg[5]
	create_data()
end


